#include <stdexcept>
#include <thread>
#include <iostream>
#include "io_pool.h"

io_context_pool::io_context_pool(std::size_t pool_size)
        : next_io_context_(0) {
    if (pool_size == 0)
        throw std::runtime_error("io_context_pool size is 0");

    // Give all the io_contexts work to do so that their run() functions will not
    // exit until they are explicitly stopped.
    for (std::size_t i = 0; i < pool_size; ++i) {
        io_contexts_.push_back(std::make_shared<boost::asio::io_context>());
    }
}

void io_context_pool::run() {
    // Create a pool of threads to run all of the io_contexts.
    std::vector<std::shared_ptr<std::thread> > threads;
    for (std::size_t i = 0; i < io_contexts_.size(); ++i) {
        std::shared_ptr<std::thread> thread(new std::thread(
                [i, this]() {
                    std::cout << "Thread " << i << " is starting.\n";
                    boost::asio::executor_work_guard<boost::asio::io_context::executor_type> work_guard =
                            boost::asio::make_work_guard(*io_contexts_[i].get());
                    work_ptr work(&work_guard);
                    work_.push_back(work);
                    io_contexts_[i]->run();
                    std::cout << "Thread " << i << " has finished.\n";
                }));
        threads.push_back(thread);
    }

    // Wait for all threads in the pool to exit.
    for (std::size_t i = 0; i < threads.size(); ++i)
        threads[i]->join();
}

void io_context_pool::stop() {
    // Explicitly stop all io_contexts.
    for (std::size_t i = 0; i < work_.size(); ++i)
        work_[i].reset();
}

boost::asio::io_context &io_context_pool::get_io_context() {
    // Use a round-robin scheme to choose the next io_context to use.
    auto &ioContext = *io_contexts_[next_io_context_];
    next_io_context_++;
    if (next_io_context_ == io_contexts_.size()) {
        next_io_context_ = 0;
    }
    return ioContext;
}

boost::asio::io_context &io_context_pool::get_io_context(size_t index) {
    index = index % io_contexts_.size();
    return *io_contexts_[index];;
}
